package com.util;

import com.consts.Const;
import org.apache.commons.codec.binary.Base64;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.rendering.PDFRenderer;
import org.apache.pdfbox.text.PDFTextStripper;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.LinkedHashMap;
import java.util.Map;

/**
 * @author Smile
 */
public class PdfUtil {


    /**
     * PDF转图片
     *
     * @param pdfBytes
     * @return
     * @throws IOException
     */
    public static Map<Integer, String> Pdf2Image(byte[] pdfBytes) throws IOException {
        // 1dpi≈11.7
        // float dpi = (maxWidth > maxHeight ? maxHeight : maxWidth) / 11.7f;

        try (PDDocument pdDocument = PDDocument.load(pdfBytes)) {
            // 文档总页数
            int pageSize = pdDocument.getNumberOfPages();
            PDFRenderer renderer = new PDFRenderer(pdDocument);
            BufferedImage image;

            Map<Integer, String> images = new LinkedHashMap<>();

            for (int i = 0; i < pageSize; i++) {
//                images = new HashMap<>();
                image = renderer.renderImage(i);
                try (ByteArrayOutputStream os = new ByteArrayOutputStream()) {
                    ImageIO.write(image, "png", os);
                    byte[] imgBytes = os.toByteArray();
                    images.put(i, imgBytes.length > 0 ? Const.img_prefix + Base64.encodeBase64String(imgBytes) : null);
                }
            }
            return images;
        }
    }

    /**
     * 提取指定PDF页文本
     *
     * @param pdfBytes
     * @param
     * @return
     * @throws Exception
     */
    public static String getTextFromPdf(byte[] pdfBytes) throws Exception {
        try (PDDocument pdfDocument = PDDocument.load(pdfBytes)) {
            return getTextFromPdf(pdfDocument, 0, pdfDocument.getNumberOfPages());
        }
    }

    public static String getTextFromPdf(byte[] pdfBytes, int index) throws Exception {
        try (PDDocument pdfDocument = PDDocument.load(pdfBytes)) {
            return getTextFromPdf(pdfDocument, index, index);
        }
    }

    public static String getTextFromPdf(byte[] pdfBytes, int start, int end) throws Exception {
        try (PDDocument pdfDocument = PDDocument.load(pdfBytes)) {
            return getTextFromPdf(pdfDocument, start, end);
        }
    }

    private static String getTextFromPdf(PDDocument pdfDocument, int start, int end) throws Exception {
        PDFTextStripper stripper = new PDFTextStripper();
        stripper.setStartPage(start);
        stripper.setEndPage(end);

        return stripper.getText(pdfDocument);
    }
}
